import cv2
import face_recognition
import pickle
import numpy as np
import os

def face_unlock():
    # 检查数据文件
    if not os.path.exists('face_data/user_face.pkl'):
        print("× 错误：未找到人脸数据！")
        print("× 请先运行 train_face.py 进行人脸录入")
        return
    
    # 加载训练数据
    with open('face_data/user_face.pkl', 'rb') as f:
        known_encodings = pickle.load(f)
    
    print("=== 人脸识别解锁 ===")
    print(f"✓ 已加载 {len(known_encodings)} 个人脸样本")
    print("✓ 请看向摄像头进行身份验证")
    print("✓ 按 Q 键退出程序")
    print()
    
    cap = cv2.VideoCapture(0)
    threshold = 0.4  # 识别阈值，越小越严格
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
            
        # 镜像翻转
        frame = cv2.flip(frame, 1)
        
        # 检测人脸
        face_locations = face_recognition.face_locations(frame)
        face_encodings = face_recognition.face_encodings(frame, face_locations)
        
        # 识别结果
        if len(face_encodings) > 0:
            # 计算相似度
            distances = face_recognition.face_distance(known_encodings, face_encodings[0])
            min_distance = np.min(distances)
            
            # 判断是否匹配
            if min_distance < threshold:
                status = "✓ UNLOCK SUCCESS!"
                color = (0, 255, 0)  # 绿色
                print("🔓 解锁成功！")  # 终端提示
            else:
                status = "× ACCESS DENIED"
                color = (0, 0, 255)  # 红色
            
            # 显示结果
            cv2.putText(frame, status, (10, 40), 
                       cv2.FONT_HERSHEY_SIMPLEX, 1.2, color, 3)
            cv2.putText(frame, f"Distance: {min_distance:.3f}", (10, 80), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
            cv2.putText(frame, f"Threshold: {threshold}", (10, 110), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
        else:
            cv2.putText(frame, "No Face Detected", (10, 40), 
                       cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 0), 2)
        
        # 画人脸框
        for (top, right, bottom, left) in face_locations:
            if len(face_encodings) > 0:
                distances = face_recognition.face_distance(known_encodings, face_encodings[0])
                min_distance = np.min(distances)
                frame_color = (0, 255, 0) if min_distance < threshold else (0, 0, 255)
            else:
                frame_color = (255, 0, 0)
            cv2.rectangle(frame, (left, top), (right, bottom), frame_color, 2)
        
        # 操作提示
        cv2.putText(frame, "Press Q to quit", (10, frame.shape[0] - 20), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
        
        cv2.imshow('Face Unlock System', frame)
        
        # 退出检测
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    cap.release()
    cv2.destroyAllWindows()
    print("程序已退出")

if __name__ == "__main__":
    face_unlock()